//! Lambda capture analysis: populate `LambdaExpr.captures` with referenced outer variables.

use crate::ast::*;
use crate::error::Result;
use std::collections::HashMap;

pub fn analyze_captures(program: &mut Program) -> Result<()> {
    let mut analyzer = CaptureAnalyzer::new();
    for decl in &mut program.declarations {
        analyzer.visit_declaration(decl)?;
    }
    Ok(())
}

struct CaptureAnalyzer {
    scopes: Vec<HashMap<String, Type>>,
}

impl CaptureAnalyzer {
    fn new() -> Self {
        Self {
            scopes: vec![HashMap::new()],
        }
    }

    fn visit_declaration(&mut self, decl: &mut Declaration) -> Result<()> {
        match decl {
            Declaration::Function(func) => self.visit_function(func),
            Declaration::Struct(struct_decl) => {
                for method in &mut struct_decl.methods {
                    self.visit_function(method)?;
                }
                Ok(())
            }
            Declaration::Impl(impl_decl) => {
                for method in &mut impl_decl.methods {
                    self.visit_function(method)?;
                }
                Ok(())
            }
            Declaration::TraitImpl(trait_impl) => {
                for method in &mut trait_impl.methods {
                    self.visit_function(method)?;
                }
                Ok(())
            }
            _ => Ok(()),
        }
    }

    fn visit_function(&mut self, func: &mut FunctionDecl) -> Result<()> {
        if let Some(body) = &mut func.body {
            self.enter_scope();
            for param in &func.parameters {
                self.insert_local(&param.name, param.param_type.clone());
            }
            self.visit_block(body)?;
            self.exit_scope();
        }
        Ok(())
    }

    fn visit_block(&mut self, block: &mut Block) -> Result<()> {
        self.enter_scope();
        for stmt in &mut block.statements {
            self.visit_statement(stmt)?;
        }
        self.exit_scope();
        Ok(())
    }

    fn visit_statement(&mut self, stmt: &mut Statement) -> Result<()> {
        match stmt {
            Statement::Expression(expr) => self.visit_expression(expr)?,
            Statement::VariableDecl(var) => {
                if let Some(init) = &mut var.initializer {
                    self.visit_expression(init)?;
                }
                let ty = var
                    .var_type
                    .clone()
                    .unwrap_or(Type::Basic(crate::ast::types::BasicType::Any));
                self.insert_local(&var.name, ty);
            }
            Statement::ConstantDecl(const_decl) => {
                self.visit_expression(&mut const_decl.value)?;
                let ty = const_decl
                    .const_type
                    .clone()
                    .unwrap_or(Type::Basic(crate::ast::types::BasicType::Any));
                self.insert_local(&const_decl.name, ty);
            }
            Statement::Assignment(assign) => {
                self.visit_expression(&mut assign.value)?;
            }
            Statement::Return(ret) => {
                if let Some(value) = &mut ret.value {
                    self.visit_expression(value)?;
                }
            }
            Statement::Block(block_stmt) => {
                self.enter_scope();
                for inner in &mut block_stmt.statements {
                    self.visit_statement(inner)?;
                }
                self.exit_scope();
            }
            Statement::If(if_stmt) => {
                self.visit_expression(&mut if_stmt.condition)?;
                self.visit_statement(&mut if_stmt.then_branch)?;
                if let Some(else_branch) = &mut if_stmt.else_branch {
                    self.visit_statement(else_branch)?;
                }
            }
            Statement::Match(match_stmt) => {
                self.visit_expression(&mut match_stmt.expr)?;
                for arm in &mut match_stmt.arms {
                    if let Some(guard) = &mut arm.guard {
                        self.visit_expression(guard)?;
                    }
                    self.visit_statement(&mut arm.body)?;
                }
            }
            Statement::While(while_stmt) => {
                self.visit_expression(&mut while_stmt.condition)?;
                self.visit_statement(&mut while_stmt.body)?;
            }
            Statement::For(for_stmt) => {
                match &mut for_stmt.variant {
                    ForVariant::Traditional {
                        initializer,
                        condition,
                        increment,
                    } => {
                        if let Some(init) = initializer {
                            self.visit_statement(init)?;
                        }
                        if let Some(cond) = condition {
                            self.visit_expression(cond)?;
                        }
                        if let Some(inc) = increment {
                            self.visit_statement(inc)?;
                        }
                    }
                    ForVariant::Condition { condition } => {
                        self.visit_expression(condition)?;
                    }
                    ForVariant::Range {
                        variable,
                        variable_type,
                        iterable,
                    } => {
                        self.visit_expression(iterable)?;
                        self.enter_scope();
                        self.insert_local(variable, variable_type.clone());
                        self.visit_statement(&mut for_stmt.body)?;
                        self.exit_scope();
                        return Ok(());
                    }
                }
                self.visit_statement(&mut for_stmt.body)?;
            }
            Statement::Defer(defer_stmt) => self.visit_expression(&mut defer_stmt.expr)?,
            Statement::Go(go_stmt) => {
                let mut call_expr = Expression::Call(go_stmt.call.clone());
                self.visit_expression(&mut call_expr)?
            }
            Statement::Select(select_stmt) => {
                for case in &mut select_stmt.cases {
                    self.visit_expression(&mut case.channel)?;
                    self.visit_statement(&mut case.body)?;
                }
                if let Some(default_case) = &mut select_stmt.default_case {
                    self.visit_statement(default_case)?;
                }
            }
            Statement::Recover(recover_stmt) => {
                for arm in &mut recover_stmt.arms {
                    if let Some(guard) = &mut arm.guard {
                        self.visit_expression(guard)?;
                    }
                    self.visit_statement(&mut arm.body)?;
                }
            }
            Statement::Export(_)
            | Statement::Break(_)
            | Statement::Continue(_)
            | Statement::Empty => {}
        }
        Ok(())
    }

    fn visit_expression(&mut self, expr: &mut Expression) -> Result<()> {
        match expr {
            Expression::Lambda(lambda) => {
                let captures = self.collect_lambda(lambda);
                lambda.captures = captures;

                self.enter_scope();
                for param in &lambda.parameters {
                    self.insert_local(&param.name, param.param_type.clone());
                }
                self.visit_expression(&mut lambda.body)?;
                self.exit_scope();
            }
            Expression::FunctionLiteral(func_lit) => {
                if let Some(lambda) = func_lit.lambda.as_mut() {
                    let captures = self.collect_lambda(lambda);
                    lambda.captures = captures;

                    self.enter_scope();
                    for param in &lambda.parameters {
                        self.insert_local(&param.name, param.param_type.clone());
                    }
                    self.visit_expression(&mut lambda.body)?;
                    self.exit_scope();
                }
            }
            Expression::Call(call) => {
                self.visit_expression(&mut call.callee)?;
                for arg in &mut call.arguments {
                    self.visit_expression(arg)?;
                }
            }
            Expression::MethodCall(method_call) => {
                self.visit_expression(&mut method_call.object)?;
                for arg in &mut method_call.arguments {
                    self.visit_expression(arg)?;
                }
            }
            Expression::Binary(binary) => {
                self.visit_expression(&mut binary.left)?;
                self.visit_expression(&mut binary.right)?;
            }
            Expression::Unary(unary) => self.visit_expression(&mut unary.operand)?,
            Expression::FieldAccess(field_access) => {
                self.visit_expression(&mut field_access.object)?
            }
            Expression::IndexAccess(index_access) => {
                self.visit_expression(&mut index_access.object)?;
                self.visit_expression(&mut index_access.index)?;
            }
            Expression::Match(match_expr) => {
                self.visit_expression(&mut match_expr.expr)?;
                for arm in &mut match_expr.arms {
                    if let Some(guard) = &mut arm.guard {
                        self.visit_expression(guard)?;
                    }
                    self.visit_expression(&mut arm.body)?;
                }
            }
            Expression::If(if_expr) => {
                self.visit_expression(&mut if_expr.condition)?;
                self.visit_expression(&mut if_expr.then_branch)?;
                if let Some(else_branch) = &mut if_expr.else_branch {
                    self.visit_expression(else_branch)?;
                }
            }
            Expression::Block(block_expr) => {
                self.enter_scope();
                for stmt in &mut block_expr.statements {
                    self.visit_statement(stmt)?;
                }
                if let Some(final_expr) = &mut block_expr.final_expr {
                    self.visit_expression(final_expr)?;
                }
                self.exit_scope();
            }
            Expression::Struct(struct_expr) => {
                for field in &mut struct_expr.fields {
                    self.visit_expression(&mut field.value)?;
                }
            }
            Expression::Array(array_expr) => {
                for element in &mut array_expr.elements {
                    self.visit_expression(element)?;
                }
            }
            Expression::Slice(slice_expr) => {
                for element in &mut slice_expr.elements {
                    self.visit_expression(element)?;
                }
            }
            Expression::Map(map_expr) => {
                for entry in &mut map_expr.entries {
                    self.visit_expression(&mut entry.key)?;
                    self.visit_expression(&mut entry.value)?;
                }
            }
            Expression::Variable(_) | Expression::Literal(_) => {}
            _ => {}
        }
        Ok(())
    }

    fn collect_lambda(&self, lambda: &LambdaExpr) -> Vec<CapturedVar> {
        let mut collector = LambdaCollector::new(&self.scopes);
        collector.push_scope();
        for param in &lambda.parameters {
            collector.insert_local(&param.name, param.param_type.clone());
        }
        collector.visit_expression(&lambda.body);
        collector.finish()
    }

    fn insert_local(&mut self, name: &str, ty: Type) {
        if let Some(scope) = self.scopes.last_mut() {
            scope.insert(name.to_string(), ty);
        }
    }

    fn enter_scope(&mut self) {
        self.scopes.push(HashMap::new());
    }

    fn exit_scope(&mut self) {
        self.scopes.pop();
    }
}

struct LambdaCollector {
    outer_scopes: Vec<HashMap<String, Type>>,
    local_scopes: Vec<HashMap<String, Type>>,
    captures: Vec<CapturedVar>,
}

impl LambdaCollector {
    fn new(outer: &[HashMap<String, Type>]) -> Self {
        Self {
            outer_scopes: outer.to_vec(),
            local_scopes: Vec::new(),
            captures: Vec::new(),
        }
    }

    fn push_scope(&mut self) {
        self.local_scopes.push(HashMap::new());
    }

    fn pop_scope(&mut self) {
        self.local_scopes.pop();
    }

    fn insert_local(&mut self, name: &str, ty: Type) {
        if let Some(scope) = self.local_scopes.last_mut() {
            scope.insert(name.to_string(), ty);
        }
    }

    fn visit_expression(&mut self, expr: &Expression) {
        match expr {
            Expression::Variable(name) => {
                if !self.is_local(name) {
                    if let Some(var_type) = self.lookup_outer(name) {
                        self.record_capture(name, var_type);
                    }
                }
            }
            Expression::Lambda(_inner) => {
                // Nested lambda handled by main analyzer
            }
            Expression::Call(call) => {
                self.visit_expression(&call.callee);
                for arg in &call.arguments {
                    self.visit_expression(arg);
                }
            }
            Expression::MethodCall(method_call) => {
                self.visit_expression(&method_call.object);
                for arg in &method_call.arguments {
                    self.visit_expression(arg);
                }
            }
            Expression::Binary(binary) => {
                self.visit_expression(&binary.left);
                self.visit_expression(&binary.right);
            }
            Expression::Unary(unary) => self.visit_expression(&unary.operand),
            Expression::FieldAccess(field_access) => self.visit_expression(&field_access.object),
            Expression::IndexAccess(index_access) => {
                self.visit_expression(&index_access.object);
                self.visit_expression(&index_access.index);
            }
            Expression::Match(match_expr) => {
                self.visit_expression(&match_expr.expr);
                for arm in &match_expr.arms {
                    if let Some(guard) = &arm.guard {
                        self.visit_expression(guard);
                    }
                    self.visit_expression(&arm.body);
                }
            }
            Expression::If(if_expr) => {
                self.visit_expression(&if_expr.condition);
                self.visit_expression(&if_expr.then_branch);
                if let Some(else_branch) = &if_expr.else_branch {
                    self.visit_expression(else_branch);
                }
            }
            Expression::Block(block_expr) => {
                self.push_scope();
                for stmt in &block_expr.statements {
                    self.visit_statement(stmt);
                }
                if let Some(final_expr) = &block_expr.final_expr {
                    self.visit_expression(final_expr);
                }
                self.pop_scope();
            }
            Expression::Struct(struct_expr) => {
                for field in &struct_expr.fields {
                    self.visit_expression(&field.value);
                }
            }
            Expression::Array(array_expr) => {
                for element in &array_expr.elements {
                    self.visit_expression(element);
                }
            }
            Expression::Slice(slice_expr) => {
                for element in &slice_expr.elements {
                    self.visit_expression(element);
                }
            }
            Expression::Map(map_expr) => {
                for entry in &map_expr.entries {
                    self.visit_expression(&entry.key);
                    self.visit_expression(&entry.value);
                }
            }
            Expression::FunctionLiteral(_) => {}
            _ => {}
        }
    }

    fn visit_statement(&mut self, stmt: &Statement) {
        match stmt {
            Statement::Expression(expr) => self.visit_expression(expr),
            Statement::VariableDecl(var) => {
                if let Some(init) = &var.initializer {
                    self.visit_expression(init);
                }
                let ty = var
                    .var_type
                    .clone()
                    .unwrap_or(Type::Basic(crate::ast::types::BasicType::Any));
                self.insert_local(&var.name, ty);
            }
            Statement::ConstantDecl(const_decl) => {
                self.visit_expression(&const_decl.value);
                let ty = const_decl
                    .const_type
                    .clone()
                    .unwrap_or(Type::Basic(crate::ast::types::BasicType::Any));
                self.insert_local(&const_decl.name, ty);
            }
            Statement::Assignment(assign) => {
                self.visit_expression(&assign.value);
                self.visit_assignment_target(&assign.target);
            }
            Statement::Return(ret) => {
                if let Some(value) = &ret.value {
                    self.visit_expression(value);
                }
            }
            Statement::Block(block_stmt) => {
                self.push_scope();
                for inner in &block_stmt.statements {
                    self.visit_statement(inner);
                }
                self.pop_scope();
            }
            Statement::If(if_stmt) => {
                self.visit_expression(&if_stmt.condition);
                self.visit_statement(&if_stmt.then_branch);
                if let Some(else_branch) = &if_stmt.else_branch {
                    self.visit_statement(else_branch);
                }
            }
            Statement::Match(match_stmt) => {
                self.visit_expression(&match_stmt.expr);
                for arm in &match_stmt.arms {
                    if let Some(guard) = &arm.guard {
                        self.visit_expression(guard);
                    }
                    self.visit_statement(&arm.body);
                }
            }
            Statement::While(while_stmt) => {
                self.visit_expression(&while_stmt.condition);
                self.visit_statement(&while_stmt.body);
            }
            Statement::For(for_stmt) => match &for_stmt.variant {
                ForVariant::Traditional {
                    initializer,
                    condition,
                    increment,
                } => {
                    if let Some(init) = initializer {
                        self.visit_statement(init);
                    }
                    if let Some(cond) = condition {
                        self.visit_expression(cond);
                    }
                    if let Some(inc) = increment {
                        self.visit_statement(inc);
                    }
                    self.visit_statement(&for_stmt.body);
                }
                ForVariant::Condition { condition } => {
                    self.visit_expression(condition);
                    self.visit_statement(&for_stmt.body);
                }
                ForVariant::Range {
                    variable,
                    variable_type,
                    iterable,
                } => {
                    self.visit_expression(iterable);
                    self.push_scope();
                    self.insert_local(variable, variable_type.clone());
                    self.visit_statement(&for_stmt.body);
                    self.pop_scope();
                }
            },
            Statement::Defer(defer_stmt) => self.visit_expression(&defer_stmt.expr),
            Statement::Go(go_stmt) => {
                let call_expr = Expression::Call(go_stmt.call.clone());
                self.visit_expression(&call_expr);
            }
            Statement::Select(select_stmt) => {
                for case in &select_stmt.cases {
                    self.visit_expression(&case.channel);
                    self.visit_statement(&case.body);
                }
                if let Some(default_case) = &select_stmt.default_case {
                    self.visit_statement(default_case);
                }
            }
            Statement::Recover(recover_stmt) => {
                for arm in &recover_stmt.arms {
                    if let Some(guard) = &arm.guard {
                        self.visit_expression(guard);
                    }
                    self.visit_statement(&arm.body);
                }
            }
            Statement::Export(_)
            | Statement::Break(_)
            | Statement::Continue(_)
            | Statement::Empty => {}
        }
    }

    fn is_local(&self, name: &str) -> bool {
        self.local_scopes
            .iter()
            .rev()
            .any(|scope| scope.contains_key(name))
    }

    fn lookup_outer(&self, name: &str) -> Option<Type> {
        self.outer_scopes
            .iter()
            .rev()
            .find_map(|scope| scope.get(name).cloned())
    }

    fn record_capture(&mut self, name: &str, ty: Type) {
        if self.captures.iter().any(|c| c.name == name) {
            return;
        }
        self.captures.push(CapturedVar {
            name: name.to_string(),
            var_type: ty,
            mode: CaptureMode::ByValue,
        });
    }

    fn visit_assignment_target(&mut self, target: &AssignmentTarget) {
        match target {
            AssignmentTarget::Variable(_) => {}
            AssignmentTarget::FieldAccess(field) => {
                self.visit_expression(&field.object);
            }
            AssignmentTarget::IndexAccess(index) => {
                self.visit_expression(&index.object);
                self.visit_expression(&index.index);
            }
            AssignmentTarget::Deref(expr) => {
                self.visit_expression(expr);
            }
        }
    }

    fn finish(self) -> Vec<CapturedVar> {
        self.captures
    }
}
